import numpy as np 
from matplotlib import pyplot as plt
import cv2
from glob import glob
from numpy import random
from img_utils import *
import os
import pandas as pd
import cPickle
from sklearn.model_selection import train_test_split

class GlassFucker(object):
    def __init__(self, color_channel='RGB', resize=None):
        self.img = None
        self.color_channel = color_channel
        self.resize = resize

    def imread(self, img_path, mask_path):
        self.img_path = img_path
        self.mask_path = mask_path
        self.img = cv2.imread(img_path)
        self.mask = cv2.imread(mask_path)
        if self.resize:
            self.img = cv2.resize(self.img, self.resize)
            self.mask = cv2.resize(self.mask, self.resize)
        self.H, self.W = self.img.shape[0:2]

        if self.color_channel=='RGB':
            self.img = cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB)
            self.mask = cv2.cvtColor(self.mask, cv2.COLOR_BGR2RGB)

    def draw_binaryMask(self):
        self.Gmask = (self.mask[:,:,0]==255) & \
                     (self.mask[:,:,1]==0) & \
                     (self.mask[:,:,2]==0)
        self.Gmask = (self.Gmask.astype(int)*255).astype(np.uint8)


    def find_bound(self):
        self.green = (self.mask[:,:,0]==0) & \
                     (self.mask[:,:,1]==255) & \
                     (self.mask[:,:,2]==0)


        self.ymin,self.xmin = np.unravel_index(np.argmax(self.green),(self.H,self.W))
        self.ymax, self.xmax = np.unravel_index(np.argmax(self.green[::-1,::-1]),(self.H,self.W))
        self.ymax, self.xmax = self.H-self.ymax-1, self.W-self.xmax-1

    def check_crayon(self,mask):
        self.blue = (mask[:,:,0]==0) & \
                     (mask[:,:,1]==0) & \
                     (mask[:,:,2]==255)
        if self.blue.sum()>0:
            return True
        else:
            return False

    def check_white(self,img):
        self.gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        self.gray = self.gray>100
        # print self.white.sum(),0.03*self.H*self.W

        if self.gray.sum() > 100:
            return True
        else:
            return False

    def check_bbx(self,mask):
        h,w = mask.shape[:2]

        #cv2 find countour
        red = (mask[:,:,0]==255) & \
                     (mask[:,:,1]==0) & \
                     (mask[:,:,2]==0)
        red = (red.astype(int)*255).astype(np.uint8)


        if red.sum()>0:
            ys, xs = np.where(red==255)
            ymin = 1. * ys.min() / h
            ymax = 1. * ys.max() / h
            xmin = 1. * xs.min() / w
            xmax = 1. * xs.max() / w
            # cv2.rectangle(mask, (int(xmin*w), int(ymin*h)), (int(xmax*w), int(ymax*h)), (0, 255, 0), 1)
            # cv2.rectangle(red, (int(xmin*w), int(ymin*h)), (int(xmax*w), int(ymax*h)), (255), 1)
            # plt.imshow(mask)
            # plt.figure()
            # plt.imshow(red)
            # plt.show()
            return [[xmin,ymin,xmax,ymax,0]]   # 0 for class 0, i.e. defect
        else:
            return []


    def grid_patch(self, patch_size, grid_num, check=True):
        if not isinstance(patch_size, tuple):
            patch_size_w, patch_size_h = patch_size, patch_size
        else:
            patch_size_w, patch_size_h = patch_size

        if not isinstance(grid_num, tuple):
            grid_num_w, grid_num_h = grid_num, grid_num
        else:
            grid_num_w, grid_num_h = grid_num

        img_patches = np.empty((grid_num_h * grid_num_w, patch_size_h, patch_size_w,3),dtype=np.uint8)
        mask_patches = np.empty((grid_num_h * grid_num_w, patch_size_h, patch_size_w, 3), dtype=np.uint8)
        coord = np.empty((grid_num_h*grid_num_w, 6), dtype=np.int)
        bboxs = []
        select = np.empty(grid_num_h*grid_num_w, dtype=bool)

        x_endsteps = np.linspace(self.xmin+patch_size_w, self.xmax, num=grid_num_w)
        y_endsteps = np.linspace(self.ymin+patch_size_h, self.ymax, num=grid_num_h)

        for i, ystep in enumerate(y_endsteps):
            for j, xstep in enumerate(x_endsteps):
                xmax = int(xstep)
                ymax = int(ystep)
                xmin = int(xstep - patch_size_w)
                ymin = int(ystep - patch_size_h)

                patch_img = self.img[ymin:ymax, xmin:xmax, :]
                patch_mask = self.mask[ymin:ymax, xmin:xmax, :]
                img_patches[i * grid_num_h+j, :, :, :] = patch_img
                mask_patches[i * grid_num_h + j, :, :, :] = patch_mask

                select[i * grid_num_h+j] = True
                if check:
                    hasCrayon = self.check_crayon(patch_mask)
                    hasWhite = self.check_white(patch_img)
                    hasBbx = self.check_bbx(patch_mask)  #[[xmin1,ymin1,xmax1,ymax1],
                                                        #  [xmin2,ymin2,xmax2,ymax2], ...]
                    if (hasCrayon or not hasWhite) or not hasBbx:
                        select[i * grid_num_h + j] = False
                    else:
                        # print hasBbx
                        bboxs.append(hasBbx)
                else:
                    hasBbx = self.check_bbx(patch_mask)
                    bboxs.append(hasBbx)



                coord[i * grid_num_h + j] = [xmin, ymin, xmax, ymax, xmax-xmin, ymax-ymin]

                # plt.figure()
                # plt.imshow(patch_img)
                # cv2.rectangle(self.img, (xmin,ymin),(xmax,ymax),(55,255,155),5)
                # plt.figure()
                # plt.imshow(self.img)
                # plt.show()
        return img_patches[select], mask_patches[select], coord[select], bboxs

    # def random_patch(self, patch_size, scale_range=(0.8,1.2),
    #                  ratio=(3./4,4./3), resize=None, with_mask=True):
    #     if not isinstance(patch_size, tuple):
    #         patch_size_w, patch_size_h = patch_size, patch_size
    #     else:
    #         patch_size_w, patch_size_h = patch_size
    #
    #     if resize is None:
    #         resize = patch_size
    #     if not isinstance(resize, tuple):
    #         resize_w, resize_h = resize, resize
    #     else:
    #         resize_w, resize_h = resize
    #
    #     for attempt in range(10):
    #         aspect_ratio = random.uniform(ratio[0], ratio[1])
    #         scale = random.uniform(scale_range[0], scale_range[1])
    #
    #         new_w = int(round(patch_size_w * aspect_ratio * scale))
    #         new_h = int(round(patch_size_h / aspect_ratio * scale))
    #
    #         if random.random() < 0.5:
    #             new_h, new_w = new_w, new_h
    #
    #         if new_w < self.xmax-self.xmin and new_h < self.ymax-self.ymin:
    #             x0 = random.randint(self.xmin, self.xmax-new_w)
    #             y0 = random.randint(self.ymin, self.ymax-new_h)
    #             out_img = fixed_crop(self.img, x0, y0, new_w, new_h, (resize_w, resize_h))
    #             if with_mask:
    #                 out_mask = fixed_crop(self.mask, x0, y0, new_w, new_h, (resize_w, resize_h))
    #             # plt.figure()
    #             # plt.imshow(out_mask)
    #             # plt.show()
    #                 if self.check_crayon(out_mask):
    #                     continue
    #             else:
    #                 out_mask = None
    #             # cv2.rectangle(self.img, (x0, y0), (x0+new_w, y0+new_h), (55, 255, 155), 5)
    #             return out_img, out_mask
    #     # Fallback
    #     out_img = fixed_crop(self.img, self.xmin, self.ymin,
    #                       self.xmax-self.xmin, self.ymax-self.ymin,
    #                       (resize_w, resize_h))
    #     if with_mask:
    #         out_mask = fixed_crop(self.mask, self.xmin, self.ymin,
    #                    self.xmax - self.xmin, self.ymax - self.ymin,
    #                    (resize_w, resize_h))
    #     return out_img, out_mask


if __name__ == '__main__':
    img_save_dir = '/media/hszc/data1/glass_data/patches/mask_no_crayon/imgs'
    train_anno_save_path = '/media/hszc/data1/glass_data/patches/mask_no_crayon/train_annos.pkl'
    val_anno_save_path = '/media/hszc/data1/glass_data/patches/mask_no_crayon/val_annos.pkl'
    if not os.path.exists(img_save_dir):
        os.makedirs(img_save_dir)


    mask_paths = sorted(glob('/media/hszc/data1/glass_data/mask_no_crayon/*.jpg'))
    mask_paths_train, mask_paths_val = train_test_split(mask_paths, test_size=0.2)
    print len(mask_paths_train), len(mask_paths_val)






    annos = {'img_pathes': [],
             'patch_pos': [],
             'bboxs_pos': []}
    GF = GlassFucker(resize=None)
    c = 0
    for mask_path in mask_paths_train:
        img_path = '/'.join(mask_path.split('/')[:-2]) + '/img_no_crayon/' + mask_path.split('/')[-1].replace('jpg','bmp')
        if not os.path.exists(img_path):
            continue
        print img_path

        GF.imread(img_path, mask_path)
        GF.find_bound()

        patches, mask_patches, coords, bboxs = GF.grid_patch(patch_size=500, grid_num=(10, 10), check=True)


        for i in xrange(len(patches)):
            img_name = mask_path.split('/')[-1].split('.')[0]+'-%d.jpg'%i
            save_path = os.path.join(img_save_dir, img_name)
            print save_path
            cv2.imwrite(save_path,patches[i])
            annos['img_pathes'].append(save_path)
            annos['patch_pos'].append(coords[i])
            annos['bboxs_pos'].append(bboxs[i])

        #     # for visual
        #     if len(bboxs[i])>0:
        #         c+=1
        #
        #         plt.figure()
        #         plt.subplot(121)
        #         plt.imshow(patches[i])
        #         for j in range(len(bboxs[i])):
        #             xmin, ymin, xmax, ymax = bboxs[i][j]
        #             xmin, ymin, xmax, ymax = int(xmin*coords[i][-2]), int(ymin*coords[i][-1]), int(xmax*coords[i][-2]), int(ymax*coords[i][-1])
        #             print xmin, ymin, xmax, ymax
        #             cv2.rectangle(patches[i], (xmin, ymin),
        #                           (xmax, ymax), (255,0,0), 2)
        #             cv2.rectangle(GF.img, (xmin+coords[i][0], ymin+coords[i][1]),
        #                           (xmax+coords[i][0], ymax+coords[i][1]), (255, 0, 0), 20)
        #
        #
        #             plt.subplot(122)
        #             plt.imshow(patches[i])
        #             plt.show()
        #
        # plt.figure()
        # plt.imshow(GF.img)
        # plt.show()
        # break

    cPickle.dump(annos, open(train_anno_save_path,'wb'))
    annos = pd.DataFrame(annos)

    print annos



    annos = {'img_pathes': [],
             'patch_pos': [],
             'bboxs_pos': []}
    GF = GlassFucker(resize=None)
    for mask_path in mask_paths_val:
        img_path = '/'.join(mask_path.split('/')[:-2]) + '/img_no_crayon/' + mask_path.split('/')[-1].replace('jpg','bmp')
        if not os.path.exists(img_path):
            continue
        print img_path

        GF.imread(img_path, mask_path)
        GF.find_bound()

        patches, mask_patches, coords, bboxs = GF.grid_patch(patch_size=500, grid_num=(10, 10), check=False)


        for i in xrange(len(patches)):
            img_name = mask_path.split('/')[-1].split('.')[0]+'-%d.jpg'%i
            save_path = os.path.join(img_save_dir, img_name)
            print save_path
            cv2.imwrite(save_path,patches[i])
            annos['img_pathes'].append(save_path)
            annos['patch_pos'].append(coords[i])
            annos['bboxs_pos'].append(bboxs[i])

    cPickle.dump(annos, open(val_anno_save_path,'wb'))
    annos = pd.DataFrame(annos)
